# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import asyncio
import logging

from absl import app, flags
from tornado.ioloop import IOLoop

from coin.util.queue.publisher import AsyncKafkaPublisher
from coin.util.queue.config import KafkaConfig

FLAGS = flags.FLAGS

flags.DEFINE_string('kafka_config_filename', None, 'kafka config')

flags.DEFINE_string('topic', 'example_topic', '')

flags.DEFINE_integer('count', 1, '')

flags.DEFINE_float('cooldown', 1, '')


async def publish():
  kafka_config = KafkaConfig.from_cmd_config(FLAGS.kafka_config_filename)
  publisher = AsyncKafkaPublisher(kafka_config=kafka_config)
  publisher.start()

  topics = FLAGS.topic.split(',')
  logging.info(topics)

  try:
    for counter in range(FLAGS.count):
      for topic in topics:
        logging.info('topic=%s count=%d', topic, counter)
        publisher.publish(topic, b'')  # str(counter).encode('ascii'))
        await asyncio.sleep(FLAGS.cooldown)
  finally:
    await asyncio.sleep(1.0)
    publisher.stop()

  ioloop = IOLoop.current()
  ioloop.stop()


def main(argv):
  ioloop = IOLoop.current()
  ioloop.add_callback(publish)
  try:
    ioloop.start()
  except KeyboardInterrupt:
    print()
  ioloop.stop()


if __name__ == "__main__":
  app.run(main)
